clc;
clear;

%% main主脚本
tic;
startTime = cputime;
pde = process_pde(4);
vector_h = [1/4];
iter_L = 9;
basis_type = "P2";

% process_showbasis(basis_type, "x");

for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh(pde, h, basis_type);

    % process_showmesh(mesh);

    O0 = sparse(mesh.Nb, mesh.Nb);
    O1 = sparse(mesh.Np, mesh.Np);
    O2 = sparse(mesh.Nb, mesh.Np);
    O3 = zeros(mesh.Np, 1);
    A1 = assemble_stiffness_matrix(mesh, mesh.basis_type, "dx", mesh.basis_type, "dx", pde.nu);
    A2 = assemble_stiffness_matrix(mesh, mesh.basis_type, "dy", mesh.basis_type, "dy", pde.nu);
    A3 = assemble_stiffness_matrix(mesh, mesh.basis_type, "dx", mesh.basis_type, "dy", pde.nu);
    A5 = assemble_stiffness_matrix(mesh, "P1", "x", mesh.basis_type, "dx", -1);
    A6 = assemble_stiffness_matrix(mesh, "P1", "x", mesh.basis_type, "dy", -1);
    A0 = assemble_stiffness_matrix(mesh, "P1", "x", "P1", "x", -1e-8);
    A = [A1+A2, O0, A5;
        O0, A1+A2, A6;
        A5', A6', A0];
    b = zeros(2*mesh.Nb+mesh.Np, 1);

    U_old = zeros(2*mesh.Nb+mesh.Np, 1);
    for l = 1:iter_L
        AN1 = assemble_nonlinear_matrix(mesh, "u1", "dx", "P2", "x", "P2", "x", U_old);
        AN2 = assemble_nonlinear_matrix(mesh, "u1", "x", "P2", "dx", "P2", "x", U_old);
        AN3 = assemble_nonlinear_matrix(mesh, "u2", "x", "P2", "dy", "P2", "x", U_old);
        AN4 = assemble_nonlinear_matrix(mesh, "u1", "dy", "P2", "x", "P2", "x", U_old);
        AN5 = assemble_nonlinear_matrix(mesh, "u2", "dx", "P2", "x", "P2", "x", U_old);
        AN6 = assemble_nonlinear_matrix(mesh, "u2", "dy", "P2", "x", "P2", "x", U_old);
        AN = [AN1+AN2+AN3, AN4, O2;
            AN5, AN6+AN2+AN3, O2;
            O2', O2', O1];
        bN1 = assemble_nonlinear_vector(mesh, "u1", "x", "u1", "dx", "P2", "x", U_old);
        bN2 = assemble_nonlinear_vector(mesh, "u2", "x", "u1", "dy", "P2", "x", U_old);
        bN3 = assemble_nonlinear_vector(mesh, "u1", "x", "u2", "dx", "P2", "x", U_old);
        bN4 = assemble_nonlinear_vector(mesh, "u2", "x", "u2", "dy", "P2", "x", U_old);
        bN = [bN1+bN2; bN3+bN4; O3];
        Al = A + AN;
        bl = b + bN;
        [Al, bl] = process_boundary(Al, bl, mesh, pde);
        U = Al\bl;
        U_old = U;
    end

    % process_plot(mesh, pde, U(1:mesh.Nb), basis_type);
    % process_plot(mesh, pde, U(mesh.Nb+1:2*mesh.Nb), basis_type);
    % process_plot(mesh, pde, U(2*mesh.Nb+1:end), "P1");
    process_field(mesh, U);

    fprintf("已完成\t%d/%d\n", i, size(vector_h,2));
end

elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%% 以下是脚本所用到的函数 %%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% process_pde
function PDE = process_pde(index)
switch index
    case 1
        PDE.index = index;
        PDE.nu = 0.0017;
        PDE.u1 = @(x,y,k)  (k == 3) * 1 + (k == 1) * 0 + (k == 2) * 0 + (k == 4) * 0;
        PDE.u2 = @(x,y,k)  (k == 3) * 0 + (k == 1) * 0 + (k == 2) * 0 + (k == 4) * 0;
        PDE.domain = [0, 1, 0, 1];
        PDE.bdary = [0, 0, 0, 0];
    case 2
        PDE.index = index;
        PDE.nu = 0.0025;
        PDE.u1 = @(x,y,k) (k == 3) * 1 + (k == 1) * 0 + (k == 2) * 0 + (k == 4) * 0;
        PDE.u2 = @(x,y,k) (k == 3) * 0 + (k == 1) * 0 + (k == 2) * 0 + (k == 4) * -1;
        PDE.domain = [0, 1, 0, 1];
        PDE.bdary = [0, 0, 0, 0];
    case 3
        PDE.index = index;
        PDE.nu = 0.003;
        PDE.u1 = @(x,y,k) (k == 3) * 1 + (k == 1) * -1 + (k == 2) * 0 + (k == 4) * 0;
        PDE.u2 = @(x,y,k) (k == 3) * 0 + (k == 1) * 0 + (k == 2) * 1 + (k == 4) * -1;
        PDE.domain = [0, 1, 0, 1];
        PDE.bdary = [0, 0, 0, 0];
    case 4
        PDE.index = index;
        PDE.nu = 1;
        PDE.u1 = @(x,y,k) (k == 3) * 1 + (k == 1) * -1;
        PDE.u2 = @(x,y,k) (k == 3) * 0 + (k == 1) * 0;
        PDE.u1dx = @(x,y) 0;
        PDE.u1dy = @(x,y) 0;
        PDE.u2dx = @(x,y) 0;
        PDE.u2dy = @(x,y) 0;
        PDE.domain = [0, 6, 0, 1];
        PDE.bdary = [0, 1, 0, 1];
    otherwise
        error("Invalid PDE index.");
end
fprintf("PDE index: %d\n", PDE.index);
fprintf("PDE domain: xmin,xmax,ymin,ymax\n");
disp(PDE.domain);
fprintf("PDE boundary: bottom,right,top,left (0=Dirichlet,1=Neumann,2=Robin)\n");
disp(PDE.bdary);
end

%% assemble_stiffness_matrix
function matrix_A = assemble_stiffness_matrix(mesh, B1_type, D1_type, B2_type, D2_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = cor * phi1(:,j) .* phi2(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_nonlinear_matrix
function matrix_A = assemble_nonlinear_matrix(mesh, u_type, du_type, B1_type, D1_type, B2_type, D2_type, U_old)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
u_map = containers.Map(...
    {'u1', 'u2'},...
    {struct('shift', 0),...
    struct('shift', mesh.Nb)});
U = u_map(u_type);
shift = U.shift;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phiu = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du_type);
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    item = phiu * U_old(mesh.elem(n,1:3)+shift);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = item .* phi1(:,j) .* phi2(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_nonlinear_vector
function vector_b = assemble_nonlinear_vector(mesh, u1_type, du1_type, u2_type, du2_type, B_type, D_type, U_old)
u_map = containers.Map(...
    {'u1', 'u2'},...
    {struct('shift', 0),...
    struct('shift', mesh.Nb)});
U1 = u_map(u1_type);
shift1 = U1.shift;
U2 = u_map(u2_type);
shift2 = U2.shift;
vector_b = zeros(mesh.Nb,1);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du2_type);
    phi3 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B_type, D_type);
    item1 = phi1 * U_old(mesh.elem(n,1:3)+shift1);
    item2 = phi2 * U_old(mesh.elem(n,1:3)+shift2);
    for j = 1:mesh.Nlb
        f =  item1 .* item2 .* phi3(:,j);
        r = gauss_weights_loc * f;
        vector_b(mesh.elem(n,j),1) = vector_b(mesh.elem(n,j),1) + r;
    end
end
end

%% process_boundary
function [A, b] = process_boundary(A, b, mesh, pde)
shift_u2 = mesh.Nb;
% Neumann
vector_b = zeros(2*mesh.Nb+mesh.Np,1);
for k = 1:size(pde.bdary,2)
    if 1 == pde.bdary(1,k)
        bde = boundary_edges(mesh, k);
        for n = 1:size(bde,1)
            points = mesh.node(bde(n,2:3),:);
            AB = [points(2,1)-points(1,1); points(2,2)-points(1,2)];
            vector_n = [0, 1; -1, 0] * AB;
            vector_n = vector_n./norm(vector_n);
            fp1 = @(x,y) pde.nu*(pde.u1dx(x,y)*vector_n(1) + pde.u1dy(x,y)*vector_n(2));
            fp2 = @(x,y) pde.nu*(pde.u2dx(x,y)*vector_n(1) + pde.u2dy(x,y)*vector_n(2));
            nk = bde(n,1);
            E = mesh.node(mesh.elem(nk,:),:);
            V = E(1:3,:);
            [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(points, "1d");
            phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
            ffp1 = fp1(gauss_points_loc(:,1), gauss_points_loc(:,2));
            ffp2 = fp2(gauss_points_loc(:,1), gauss_points_loc(:,2));
            for i = 1:mesh.Nlb
                f1 = ffp1 .* phi(:,i);
                f2 = ffp2 .* phi(:,i);
                r1 = gauss_weights_loc * f1;
                r2 = gauss_weights_loc * f2;
                vector_b(mesh.elem(nk,i),1) = vector_b(mesh.elem(nk,i),1) + r1;
                vector_b(mesh.elem(nk,i)+shift_u2,1) = vector_b(mesh.elem(nk,i)+shift_u2,1) + r2;
            end
        end
    end
end
b = b + vector_b;
% Dirichlet
for k = 1:size(pde.bdary,2)
    if 0 == pde.bdary(1,k)
        bdn1 = boundary_nodes(mesh, k);
        A(bdn1,:) = 0;
        A(sub2ind(size(A), bdn1, bdn1)) = 1;
        b(bdn1) = pde.u1(mesh.node(bdn1,1),mesh.node(bdn1,2),k);
        bdn2 = bdn1 + shift_u2;
        A(bdn2,:) = 0;
        A(sub2ind(size(A), bdn2, bdn2)) = 1;
        b(bdn2) = pde.u2(mesh.node(bdn1,1),mesh.node(bdn1,2),k);
    end
end
end